
(*
 * Learner - performs unsupervised learning to find (near-)optimal management
 *)

open ExtList
open Lazy
open Random
open Random.State
open Printf

open Randif
open Types
open Management
open Managementstore
open Banking
open Bmap
open World2;;

type global_params = {
    checkpoint_frequency : int ;
    longoutput_frequency : int ;
    fileroot : string;
    debug : bool }

type simParams = {
    random_state : State.t;
    maxSteps:int ;
    currentSteps: int;
    retentionBound:int;
    mutator:mutator;
    bestManagement: management list; 
    mgtStore:managementStore;
    simid: int }

let buildfilename gp num =
  (Printf.sprintf "%s/checkpoint-%d.rcp" gp.fileroot num) 

let checkpoint_sim filename ( s : simParams )  =
  let oc = open_out_bin filename in
  Marshal.to_channel oc s []


let uncheckpoint_sim filename =
  let ic = open_in_bin filename in
  ( Marshal.from_channel ic : simParams )

let rec rut n l lorig =
  match n with
  | 0 -> []
  | _ -> 
      if l == []
      then (List.hd lorig) :: (rut (n - 1) (List.tl lorig) lorig)
      else (List.hd l) :: (rut (n - 1) (List.tl l) lorig)

let repeatUpTo n l = 
  rut n [] l

let construct_bmaps thresh blist =
  let assbm func ls =
    let bmaps = List.map func ls in
    List.fold_left Bmap.merge (Bmap.empty cs_max thresh ) (bmaps)
  in
  BayesM {
    accepted_lenders = assbm (fun x -> x.accepted_lender_contracts) blist;
    rejected_lenders = assbm (fun x -> x.rejected_lender_contracts) blist;
    accepted_borrowers = assbm (fun x -> x.accepted_borrower_contracts) blist;
    rejected_borrowers = assbm (fun x -> x.rejected_borrower_contracts) blist;
    bayes_dividend_rate = PNum 0.1  }
    
let assembleManagement num fromfst mstore muts =
  let storeMgt = getBestMgt (fromfst) mstore in
  let mms = List.take (num - fromfst) muts in
  shuffle (repeatUpTo num (storeMgt @ mms))

let print_managementassoc ms =
  let (scs,mgts') = (List.split ms) in
  let mgts = List.map (fun id -> (print_management id)) mgts' in
  List.fold_left2 (fun a b c -> a ^ b ^ " : " ^ c ) "" (List.map (sprintf "%8.3f") scs) mgts

let stepSim sps w globp =
  if sps.currentSteps mod globp.checkpoint_frequency == 0
  then let fname = buildfilename globp sps.currentSteps in
           checkpoint_sim fname sps
  else () ;
  if sps.currentSteps mod globp.longoutput_frequency == 0
  then  print_endline (print_table_best sps.mgtStore)
  else () ;
  Random.set_state sps.random_state;
  (* if globp.debug then (print_endline "Mutating ...") else ();
  print_endline (print_managementstore sps.mgtStore); *)
  let new_management = sps.bestManagement in
  (* let newManagement = assembleManagement 
      (List.length (w.banks)) sps.retentionBound 
      sps.mgtStore mut in *)
  let nw0 = assignManagement new_management w in
  if globp.debug then (print_endline "Starting worldstep ...") else ();
  let nw = n_step_world globp.debug nw0 in
  if globp.debug then (print_endline "Step world done; constructing mgtassoc ...") else ();
  (* let mgtAssoc = List.map (fun b -> (bank_score b,(construct_bmaps [b]))) nw.banks in *)
  let meanprof = meanv  (List.map (fun b -> (bank_score b)) nw.banks) in
  let newmgt = merge_bayes_management (construct_bmaps meanprof nw.banks) (List.hd sps.bestManagement) in  (* uses the new threshold = meanprof *)
  let mgtAssoc = List.combine (List.map (fun b -> (bank_score b)) nw.banks)
                              (replicate (List.length nw.banks) newmgt) in
  (* print_string (bsummary sps.simid sps.currentSteps nw);
  print_string (print_managementassoc mgtAssoc); *)
  (* Printf.printf "%s\n" (Std.dump (List.map fst sps.mgtStore.sc_to_id));*)
  if globp.debug then (print_endline "Constructing new sim state ...") else ();
  flush stdout;
  { sps with 
     random_state = Random.get_state () ;   (* no use of rng below !!! *) 
     currentSteps = (sps.currentSteps + 1);
     bestManagement = List.map snd
       (List.sort ~cmp:(fun (x,_) (y,_)  -> compare y x) mgtAssoc); 
     (* mgtStore = addList sps.mgtStore mgtAssoc *) }

let rec manyWorlds sps w gp =
  if sps.currentSteps > sps.maxSteps
  then sps
  else 
    let nsps = stepSim sps w gp in
    (manyWorlds nsps w gp)

